from simgrid import Actor, Engine, Host, Mailbox, this_actor, Mutex
import sys
import threading
from time import sleep
import socket

# master-begin
def master(*args):
  tasks = args[0]
  actions = args[1]
  #======== Create Host =======================
  hosts = Engine.instance.all_hosts
  p = 0
  for host in hosts:
    # print(host.name, host.speed)
    if host.name == 'Tremblay':
        continue
    # this_actor.info(f"Create host: {host.name}, pstate: {host.pstate}")
    Actor.create(host.name, host, worker, p)
    p += 1
  workers_cnt = len(hosts)-1 #6 workers
  workers = [Mailbox.by_name(str(i)) for i in range(workers_cnt)] 
  #========= Send Task ===============================
  for i in range(len(actions)):
    action = actions[i]
    mailbox = workers[action]
    compute_cost = tasks[i][0]
    communicate_cost = tasks[i][1]
    # this_actor.info(f"Send message to {mailbox.name} ({compute_cost}, {communicate_cost})")
    # this_actor.info(f"Send tasks: {compute_cost} to {action}")
    mailbox.put(compute_cost, communicate_cost)
  #========== Finish ================================ 
  # this_actor.info("All tasks have been dispatched. Request all workers to stop.")
  for i in range (workers_cnt):
      # The workers stop when receiving a negative compute_cost
      mailbox = workers[i] 
      # this_actor.info(f"Send tasks: -1 to {i}")
      mailbox.put(-1, 0)
# master-end

# worker-begin
def worker(*args):
  # assert len(args) == 0, "The worker expects to not get any argument"
  #argument 0, 1, 2, 3, 4
  # this_actor.info(f"worker argument: {args[0]}")
  mailbox = Mailbox.by_name(str(args[0]))
  done = False
  t = 0
  while not done:
    compute_cost = mailbox.get()
    # this_actor.info(f"get message: {compute_cost}")
    if compute_cost > 0: # If compute_cost is valid, execute a computation of that cost 
      t1 = Engine.instance.clock
      this_actor.execute(compute_cost)
      t2 = Engine.instance.clock
      t += t2 - t1
    else: # Stop when receiving an invalid compute_cost
      done = True
  fp = open('./log.txt', 'a')
  # s = str(args[0]) + ': ' + str(t2-t1) + '\n'
  s = this_actor.get_host().name + ': ' + str(t) + '\n'
  fp.write(s)
  fp.close()
  # this_actor.info(f"Exiting now. Use time: {t2-t1}")
# worker-end

if __name__ == '__main__':
  n = 9
  tasks = []
  for i in range(n):
      compute_cost = 5000000000
      communicate_cost = 100000
      tasks.append([compute_cost, communicate_cost])
  actions = [0, 2, 0, 1, 2, 3, 4, 5, 4]
  #================ SimGrid ==========================
  argv = ['app_masterworkers.py', '../platform/platform.xml', 'deployment.xml']
  e = Engine(argv)
  # self.e.register_actor("worker", worker)
  # Load the platform description and then deploy the application

  e.load_platform(argv[1]) 
  Actor.create('Tremblay', Host.by_name('Tremblay'), master, tasks, actions)
  e.run()